import pandas as pd
import numpy as np
from datetime import datetime
import csv


'''筛选沪深交易所个股 （创业板个股2020-8-24日改为20cm张跌停板）'''
def filter_20cm(df):
    df2 = df[df['code'].astype(str).str.startswith(('30'))]
    df2 = df2[df2['date'] > '2020-07-31']
    return df2

    
'''筛选沪深交易所个股 （创业板个股2020-8-24日改为20cm张跌停板）'''
def filter_10cm(df):
    df1 = df[df['code'].astype(str).str.startswith(('00','60',))]
    df2 = df[df['code'].astype(str).str.startswith(('30'))]
    df2 = df2[df2['date'] < '2020-07-31']
    # 按行拼接
    final_df = pd.concat([df1, df2], axis=0, ignore_index=True)
    return final_df

''' 输出时间：2025年02月27日 15时30分45秒  '''
def get_current_time_zh() -> str:
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

''' 格式化时间 '''
def format_date(df):
    df['date'] = pd.to_datetime(df['date']).dt.strftime('%m-%d')
    return df

''' 重置索引 '''
def r_index(df):
    return df.reset_index(drop=True)

''' Dataframe 筛选 startWith '''
def df_start_sort(df,start_word,sort_c):
    df = df[df['code'].str.startswith(start_word)]
    return sort_df(df,sort_c).reset_index(drop=True)


''' 将dict转成Dataframe '''
def dict_to_dataframe(origin_dict):
    # 创建一个空的DataFrame，用于存放最终结果
    result_df = pd.DataFrame()
    # 遍历字典，将每个DataFrame与股票代码和日期索引合并，并追加到result_df中
    for code, df in origin_dict.items():
        # 将索引（日期）和股票代码作为列添加到DataFrame中
        temp_df = df.reset_index().rename(columns={'index': 'date'})
        temp_df['code'] = code  # 添加股票代码列
        result_df = pd.concat([result_df, temp_df], ignore_index=True)  # 合并DataFrame
    return result_df
    
''' 定义一个函数，根据股票代码的开头追加后缀 '''
def append_suffix(code):
    if str(code).startswith('00') or str(code).startswith('30'):
        return code + '.SZ'
    elif str(code).startswith('60') or str(code).startswith('68'):
        return code + '.SH'
    else:
        return code + '.BJ'


''' 按 'code' 分组并获取每个分组的倒数第N天的数据 '''
# def get_last_day_df(df,n):
#     return df.groupby('code',as_index=False).nth(-n).reset_index(drop=True)
''' 按 'code' 分组并获取每个分组的倒数第N天的数据 '''
def get_last_day_df(df,n):
    df = df.assign(date=pd.to_datetime(df['date'])).sort_values('date', ascending=False).drop_duplicates('code', keep='first').reset_index(drop=True)                           
    return df

# def get_latest_date_data_memopt(df):
#     """内存优化版本，适合极大数据集"""
#     try:
#         # 直接使用日期字符串比较（避免类型转换）
#         date_strs = df['date'].astype(str)
#         max_date = date_strs.max()
        
#         # 使用布尔索引筛选
#         mask = date_strs == max_date
#         return df[mask].reset_index(drop=True)
    
#     except Exception as e:
#         print(f"获取最新日期数据失败：{str(e)}")
#         raise

import pandas as pd
from datetime import timedelta

def get_previous_days_data(df, days_before=1):
    """
    获取最新日期前 N 天的数据
    
    参数:
    df: 包含日期列的DataFrame
    days_before: 要获取的前几天数据，默认为1（前一天）
    
    返回:
    指定天数的所有行
    """
    try:
        # 1. 确保日期列是日期时间类型
        df = df.copy()
        df['date'] = pd.to_datetime(df['date'])
        
        # 2. 找到最新日期
        latest_date = df['date'].max()
        
        # 3. 计算目标日期
        target_date = latest_date - timedelta(days=days_before)
        
        # 4. 筛选目标日期的数据
        previous_data = df[df['date'] == target_date].copy()
        
        # 5. 重置索引并返回
        return previous_data.reset_index(drop=True)
    
    except Exception as e:
        print(f"获取前{days_before}天数据失败：{str(e)}")
        raise

def get_last_n_days_complete(df, days=10):
    """
    获取完整的最近N个日历日数据（即使某些日没有数据）
    
    返回:
    包含日期索引的DataFrame，确保每天都有记录（可能为空）
    """
    try:
        df = df.copy()
        df['date'] = pd.to_datetime(df['date'])
        
        # 找到最新日期
        latest_date = df['date'].max()
        start_date = latest_date - timedelta(days=days-1)
        
        # 创建完整的日期范围
        date_range = pd.date_range(start=start_date, end=latest_date, name='date')
        
        # 创建日期索引
        date_df = pd.DataFrame(index=date_range)
        
        # 按日期分组（保留原始数据结构）
        grouped = df.groupby('date')
        
        # 重新索引确保所有日期都存在
        result = []
        for date in date_range:
            if date in grouped.groups:
                result.append(grouped.get_group(date))
            else:
                # 创建空行保持日期完整
                empty_row = pd.DataFrame(columns=df.columns)
                empty_row['date'] = [date]
                result.append(empty_row)
                
        return pd.concat(result).reset_index(drop=True)
    
    except Exception as e:
        print(f"获取完整{days}天数据失败：{str(e)}")
        raise


        
        
'''按‘code’分组并获取最近10天的数据 '''
def get_last_days_df(df,n):
    df = df.sort_values(by='date')
    return df.groupby('code').tail(n).reset_index(drop=True)


def p_col_len(df):
    print(len(df),'\n  ',df.columns)
def p_col(df):
    print(len(df))
    
''' 定义一个函数，如果数据是数字，得到数字单位为：亿'''
def divide_by_100(x):
    try:
        return float(x) / 100000000  # 尝试将数据转换为浮点数并除以100（1亿）
    except (ValueError, TypeError):
        return x  # 如果转换失败，返回原值

''' 大单定义计算同比增幅的函数 '''
def calculate_yoy_growth(group):
    group = group.sort_values(by=['code', 'date'])
    prev_xl = group['xl'].shift(1)
    group['xl_yoy'] = np.where(prev_xl < 1, group['xl'], group['xl'] / prev_xl )
    group = group.reset_index(drop=True)
    return group

''' df排序 '''
def sort_df(df,sort_column):
    return df.sort_values(by=[sort_column],ascending=True).reset_index(drop=True)
    
def drop_code(df):
    return df.drop_duplicates(subset=['code'])

is_save_csv = True
def save_csv(df,name):
    if is_save_csv:
        df.to_csv('input/'+str(name),index=False,encoding='utf-8-sig')
    else :
        print('不用保存文件：'+name)


def to_xg(codes):
    sector_name='xg'
    # xg_df = custom_sector(sector_name)
    # custom_sector(sector_name, 'pop', xg_df['symbol']) # 移除板块1的股票
    custom_sector(sector_name, 'remove') # 删除板块1
    custom_sector(sector_name, 'insert', codes) # 新建板块1的股票列表

''' 分组存储  '''
def to_xg_split(codes):
    list_split= ['00','60','30']
    for code in list_split:
        filtered_list = [item for item in codes if str(item).startswith(code)]
        sector_name = 'xg'+code
        print(sector_name,len(filtered_list))
        custom_sector(sector_name, 'remove') # 删除板块1
        custom_sector(sector_name, 'insert', filtered_list) # 新建板块1的股票列表


""" 获取2013-2025年日线数据，每120天存储一个文件  """
def req_all_days(code_list):
    # 调用示例（包含代码过滤）
    filtered_codes = [code for code in code_list 
                 if not code.endswith('.BJ') 
                 and not code.startswith('68')]
    """获取2017-2025年日线数据，每220天存储一个文件"""
    # 初始化时间参数
    # base_date = pd.Timestamp("2013-01-01")
    base_date = pd.Timestamp("2025-01-01")
    delta = pd.Timedelta(days=120)
    final_date = pd.Timestamp("2025-03-31")
    
    # 生成时间序列
    date_sequence = []
    while base_date <= final_date:
        date_sequence.append(base_date.strftime("%Y%m%d"))
        base_date += delta

    # 遍历每个时间节点
    for i, end_date_str in enumerate(date_sequence):
        print(end_date_str)
        # 获取数据（返回字典类型）
        raw_data = get_price(
            filtered_codes,
            end_date=end_date_str,
            fre_step='1d',
            fields=['open', 'close', 'high', 'low', 'volume',
                    'quote_rate','high_limit','turnover','turnover_rate'],
            skip_paused=True,
            fq='pre',
            bar_count=120,
            is_panel=0
        )
        
        # 转换为DataFrame并检查空值
        day_df = dict_to_dataframe(raw_data)  # 假设该函数能处理空字典返回空DataFrame
        if not day_df.empty:
            day_df = day_df.round(2)
            filename = f"day_{end_date_str}_batch{i+1}.csv"
            save_csv(day_df, filename)
            print(f"已存储：{filename}（数据量：{len(day_df)}行）")
        else:
            print(f"无有效数据：{end_date_str}")

    return "数据获取完成"


# req_days(filtered_codes)


'''  切分文件,切分Dataframe对象 '''
def split_dataframe(df, num_files):
    chunks = np.array_split(df, num_files)
    for i, chunk in enumerate(chunks):
        chunk.to_csv('input/part2_'+str(i+1)+'.csv', index=False)

''' 切分源文件 '''
def split_csv(input_file, output_prefix, num_chunks):
    with open(input_file, 'r') as f:
        reader = csv.reader(f)
        header = next(reader)
        total_rows = sum(1 for _ in reader)
        f.seek(0)
        next(reader)
        chunk_size = (total_rows + num_chunks - 1) // num_chunks
        writers = [csv.writer(open(f'{output_prefix}_{i}.csv', 'w', newline='')) for i in range(num_chunks)]
        [w.writerow(header) for w in writers]
        for i, row in enumerate(reader):
            writers[i//chunk_size].writerow(row)
        [w.__dict__['writer'].close() for w in writers]


import pandas as pd

def split_and_save(df: pd.DataFrame, 
                  output_dir: str = "output",
                  prefix: str = "part") -> None:
    """
    将DataFrame等分3份存储到指定目录
    
    参数:
        df: 要分割的DataFrame
        output_dir: 输出目录路径（默认output）
        prefix: 文件名前缀（默认part）
    
    特性:
        1. 自动处理余数分配[1](@ref)
        2. 直接使用字符串路径拼接
        3. 依赖目录已存在（需手动创建output_dir）
    """
    total = len(df)
    base, rem = divmod(total, 3)
    
    # 生成分片索引（网页1的iloc分割法）
    slices = [
        (0, base + rem),
        (base + rem, base*2 + rem),
        (base*2 + rem, total)
    ]
    
    # 直接存储CSV文件（网页2的批量存储逻辑）
    for i, (s, e) in enumerate(slices, 1):
        df.iloc[s:e].to_csv(
            f"{output_dir}/{prefix}_{i}.csv", 
            index=False
        )

def load_and_merge(input_dir: str = "input",
                  prefix: str = "part") -> pd.DataFrame:
    """
    合并指定目录的分片文件
    
    参数:
        input_dir: 输入目录路径（默认input）
        prefix: 文件名前缀（默认part）
    
    返回:
        合并后的完整DataFrame
    
    特性:
        1. 预设固定分片数[2](@ref)
        2. 直接构造文件路径
    """
    # 生成固定文件名列表
    files = [f"{input_dir}/{prefix}_{i}.csv" for i in [1,2,3]]
    
    # 合并数据（网页2的concat方法）
    return pd.concat(
        (pd.read_csv(f) for f in files),
        ignore_index=True
    )

    
from functools import reduce
def merge_df(df_list, key_columns=['date', 'code']):
    """
    根据 date 和 code 列合并多个 DataFrame
    :param df_list: 需要合并的 DataFrame 列表（如 [raise_df, ma_df, ...]）
    :param key_columns: 合并键列，默认 ['date', 'code']
    :return: 合并后的 DataFrame
    """
    # 步骤1：预处理（统一列格式）
    processed_dfs = []
    for df in df_list:
        # 强制统一日期和代码格式[1,3](@ref)
        df = df.copy()
        df['date'] = pd.to_datetime(df['date'])
        df['code'] = df['code'].astype(str).str.zfill(6)  # 股票代码标准化为6位字符串
        processed_dfs.append(df)

    # 步骤2：顺序合并（内存优化策略）
    merged_df = reduce(
        lambda left, right: pd.merge(
            left, right,
            on=key_columns,
            how='outer',  # 保留所有数据[1](@ref)
            suffixes=('', '_DROP'),  # 自动标记重复列[3](@ref)
            validate='many_to_one'  # 校验键唯一性[1](@ref)
        ),
        processed_dfs
    )

    # 步骤3：清理重复列（保留第一个出现列）
    merged_df = merged_df.loc[:, ~merged_df.columns.str.endswith('_DROP')]
    merged_df.round(2).to_csv('input/feature_v2.csv')

    return merged_df


from functools import reduce
import pandas as pd

def merge_all_df(df_list, key_columns=['date', 'code']):
    """
    修复版：安全合并多个DataFrame
    新增功能：
    1. 动态数据类型校验（网页1）
    2. 智能去重策略（网页3）
    3. 验证模式自适应（网页6）
    """
    # 步骤1：强化预处理
    processed_dfs = []
    for df in df_list:
        df = df.copy()
        
        # 统一格式（网页1建议）
        df['date'] = pd.to_datetime(df['date']).dt.normalize()  # 去除时间部分
        # df['code'] = df['code'].astype(str).str.extract('(\d{6})')[0]  # 提取纯数字代码
        df['code'] = df['code'].astype(str).str.extract(r'(\d{6}\.\w{2})')[0]  # 匹配格式如 600000.SH

        # 智能去重（网页3方案增强）
        df = df.sort_values(by='date', ascending=False)  # 按日期降序排列
        df = df.drop_duplicates(
            subset=key_columns,
            keep='first'  # 保留最新记录
        )
        
        processed_dfs.append(df)
    # print(processed_dfs)
    # 步骤2：安全合并（网页4、网页6方案结合）
    def safe_merge(left, right):
        # 动态验证模式（网页6）
        right_unique = right[key_columns].duplicated().sum() == 0
        validate_type = 'many_to_one' if right_unique else 'many_to_many'
        
        return pd.merge(
            left, right,
            on=key_columns,
            how='outer',
            suffixes=('', '_DROP'),
            validate=validate_type  # 动态验证（网页6）
        )
    
    # 分块合并优化内存（网页6建议）
    chunk_size = 2  # 每次合并2个DF
    chunks = [processed_dfs[i:i+chunk_size] 
             for i in range(0, len(processed_dfs), chunk_size)]
    
    merged_df = pd.DataFrame()
    for chunk in chunks:
        temp_df = reduce(safe_merge, chunk)
        merged_df = pd.merge(
            merged_df, temp_df, 
            on=key_columns, how='outer'
        ) if not merged_df.empty else temp_df

    # 步骤3：高级清理
    # 删除所有_DROP结尾的列
    merged_df = merged_df.loc[:, ~merged_df.columns.str.endswith('_DROP')]
    # 删除全为NaN的列
    merged_df = merged_df.dropna(axis=1, how='all')
    
    # 按日期排序后存储（网页6推荐）
    merged_df.sort_values(by='date', inplace=True)
    merged_df.round(2).to_csv('input/feature_v2.csv', index=False)
    
    return merged_df


def drop_base_columns(df):
    # df.drop(['open', 'close', 'high', 'low', 'quote_rate','high_limit', 'turnover', 't_rate','xl', 'xl_rate','gn_zt_ratio','hy_zt_ratio' ], axis=1,errors='ignore', inplace=True)

    base_drop_cols = ['open', 'close', 'high', 'low', 'quote_rate','high_limit', 'turnover', 't_rate','xl', 'xl_rate','gn_zt_ratio','hy_zt_ratio' ]
    # 安全删除（网页2建议）
    all_drop_cols = [col for col in (base_drop_cols ) if col in df.columns]
    df.drop(columns=all_drop_cols, axis=1, inplace=True, errors='ignore')
    return df



def tmp_convert_boolean_int(df):
    # 布尔字段转换
    bool_cols = ['zs_is_zt', 'zs_is_gt4', 'zs_is_gt7', 'zs_is_gt9']
    for col in bool_cols:
        df[col] = df[col].map({'True': True, 'False': False})  # 或直接 astype(bool)
    
    # 数值字段转换
    num_cols = ['zs_gt4', 'zs_gt7', 'zs_gt9']
    for col in num_cols:
        df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0).astype('int32')
    return df